/*
    This file is part of the Time-of-Flight Tracer program. It is not part of
    the original PBRT source distribution. See the included license file for
    more information.

    Copyright(c) 2014 Microsoft Corporation

    Author: Phil Pitts
*/


#include "stdafx.h"
#include "integrators/tofbipath.h"
#include "scene.h"
#include "intersection.h"
#include "paramset.h"


ToFBipathIntegrator::ToFBipathIntegrator(int md, float lambdaMin, 
	float lambdaMax) : BipathIntegrator(md), lambdaMin(lambdaMin),
	lambdaMax(lambdaMax) {
}


IntegrationResult ToFBipathIntegrator::Li(const Scene *scene,
	const Renderer *renderer, const RayDifferential &ray, 
	const Intersection &isect, const Sample *sample, RNG &rng, 
	MemoryArena &arena) const {
	// Declare common path integration variables
	Spectrum L(0.f), escapedAlpha(0.f);
	RayDifferential escapedRay;
	bool prevSpecular = true;
	PathData data(scene, renderer, rng, arena);

	float lightRayPdf;
	RayDifferential lightRay;
	Spectrum lightAlpha = SampleLight(data, sample, &lightRayPdf,
		&lightRay);
	if (lightAlpha.IsBlack() || lightRayPdf == 0.f) {
		return PathIntegrator::Li(scene, renderer, ray, isect, sample, rng, arena);
	}

	ToFPathVertex *cameraPath = new ToFPathVertex[maxDepth];
	ToFPathVertex *lightPath = new ToFPathVertex[maxDepth];
	int cameraPathLength = GeneratePath(data, ray, isect, sample, cameraPath,
		&escapedRay, &escapedAlpha);
	int lightPathLength = GenerateLightPath(data, lightRay, sample, lightAlpha,
		lightPath);

	// Compute number of specular vertices for each path length
	int nVerts = cameraPathLength + lightPathLength + 2;
	int *nSpecularVertices = new int[nVerts];
	GetNumSpecularVertices(cameraPath, lightPath, cameraPathLength,
		lightPathLength, nSpecularVertices);

	vector<HistogramSample> histSamples;
	for (int i = 0; i < cameraPathLength; ++i) {
		// Initialize basic variables for camera path vertex
		const ToFPathVertex &vc = cameraPath[i];

		// Add emitted light from vertex if appropriate
		if (prevSpecular)
			L += vc.alpha * vc.isect.Le(vc.wPrev);

		// Add direct illumination for path vertex
		float light2isect;
		Spectrum Ld = DirectLightingVertex(data, sample, i, vc, &light2isect);

		prevSpecular = vc.specularBounce;
		Ld = Ld / (i + 1 - nSpecularVertices[i + 1]);

		HistogramSample histSample;
		histSample.distance = vc.pathlength + light2isect;
		histSample.radiance = Ld.Average(lambdaMin, lambdaMax);
		histSamples.push_back(histSample);

		L += Ld;

		if (!vc.specularBounce) {
			// Loop over light path vertices and connect to camera vertex
			for (int j = 0; j < lightPathLength; ++j) {
				Spectrum lightL = LightPathL(data, *sample, lightPath, vc,
					nSpecularVertices, i, j);

				// Store ToF data for the vertex
				HistogramSample histSample;
				Point &pc = cameraPath[i].isect.dg.p;
				Point &pl = lightPath[j].isect.dg.p;
				float light2camera = (pc - pl).Length();
				histSample.distance = cameraPath[i].pathlength +
					lightPath[j].pathlength + light2camera;
				histSample.radiance = lightL.Average(lambdaMin, lambdaMax);

				L += lightL;
			}
		}
	}

	delete[] nSpecularVertices;
	delete[] cameraPath;
	delete[] lightPath;
	return IntegrationResult(L, histSamples);
}


uint32_t ToFBipathIntegrator::GeneratePath(PathData &data, const RayDifferential r,
	const Intersection &firstIsect, const Sample* sample, ToFPathVertex *path,
	RayDifferential *escapedRay, Spectrum *escapedAlpha) const {
	// Declare common path integration variables
	Spectrum pathThroughput = 1.f;
	RayDifferential ray(r);

	ToFPathVertex *v = &path[0];
	v->isect = firstIsect;

	int length = 0;
	for (; length < maxDepth; length++) {
		// Try to generate next vertex of ray path
		v = &path[length];
		if (length > 0) {
			if (!data.scene->Intersect(ray, &v->isect)) {
				*escapedAlpha = pathThroughput;
				*escapedRay = ray;
				break;
			}
		}
		SavePathVertexInfo(v, ray, pathThroughput, data.arena);
		SaveToFVertexInfo(path, length, ray);

		// Sample BSDF to get new path direction
		float pdf;
		BxDFType flags;
		BSDFSample BSDFSampleOut = BSDFSample(sample, pathSampleOffset, length);
		Spectrum f = v->bsdf->Sample_f(-ray.d, &v->wNext, BSDFSampleOut, &pdf,
			BSDF_ALL, &flags);
		v->specularBounce = (flags & BSDF_SPECULAR) != 0;
		v->nSpecularComponents = v->bsdf->NumComponents(BxDFType(BSDF_SPECULAR |
			BSDF_REFLECTION | BSDF_TRANSMISSION));

		if (f.IsBlack() || pdf == 0. || !RussianRoulette(data, *v, f, pdf, &ray,
			&pathThroughput)) {
			return length + 1;
		}
	}
	return length;
}


void ToFBipathIntegrator::SaveToFVertexInfo(ToFPathVertex *path, int pathIdx,
	const RayDifferential &ray) const {
	path[pathIdx].pathlength = (ray.o - path[pathIdx].isect.dg.p).Length();
	
	if (pathIdx > 0) {
		path[pathIdx].pathlength += path[pathIdx - 1].pathlength;
	}
}


uint32_t ToFBipathIntegrator::GenerateLightPath(PathData &data, 
	const RayDifferential lightRay, const Sample* sample, const Spectrum &a, 
	ToFPathVertex *path) const {
	// Declare common path integration variables
	RayDifferential ray(lightRay);
	Spectrum alpha(a);

	int length = 0;
	for (; length < maxDepth; length++) {
		// Try to generate next vertex of ray path
		ToFPathVertex *v = &path[length];
		if (!data.scene->Intersect(ray, &v->isect)) break;
		SavePathVertexInfo(v, ray, alpha, data.arena);
		SaveToFVertexInfo(path, length, ray);

		// Sample BSDF to get new path direction
		float pdf;
		BxDFType flags;
		BSDFSample BSDFSampleOut = BSDFSample(sample, lightPathSampleOffset,
			length);
		Spectrum f = v->bsdf->Sample_f(-ray.d, &v->wNext, BSDFSampleOut, &pdf,
			BSDF_ALL, &flags);
		v->specularBounce = (flags & BSDF_SPECULAR) != 0;
		v->nSpecularComponents = v->bsdf->NumComponents(BxDFType(BSDF_SPECULAR |
			BSDF_REFLECTION | BSDF_TRANSMISSION));

		if (f.IsBlack() || pdf == 0. || !RussianRoulette(data, *v, f, pdf, &ray,
			&alpha)) {
			return length + 1;
		}
	}
	return length;
}


Spectrum ToFBipathIntegrator::DirectLightingVertex(PathData &data,
	const Sample *sample, int pathIdx, const ToFPathVertex &v, 
	float *light2isect) const {
	Spectrum L(0.f);
	const Point &pc = v.bsdf->dgShading.p;
	const Normal &n = v.bsdf->dgShading.nn;

	float lightPdf;
	uint32_t lightNum = lightDistribution->SampleDiscrete(
		sample->oneD[lightNumOffset][pathIdx], &lightPdf);
	const Light *light = data.scene->lights[lightNum];
	BSDFSample bsdfSample(sample, bsdfSampleOffset, pathIdx);
	LightSample lightSample(sample, lightSampleOffset, pathIdx);

	L = v.alpha *
		EstimateDirect(data.scene, data.renderer, data.arena, light, pc, n, 
		v.wPrev, v.isect.rayEpsilon, sample->time, v.bsdf, data.rng, lightSample,
		bsdfSample, BxDFType(BSDF_ALL & ~BSDF_SPECULAR)) / lightPdf;

	// Extract distance from point to light for ToF histogram
	Normal normal;
	Ray lightRay;
	light->Sample_L(data.scene, lightSample, data.rng.RandomFloat(),
		data.rng.RandomFloat(), sample->time, &lightRay,
		&normal, &lightPdf);
	Point pl = lightRay.o;
	*light2isect = Distance(pl, pc);

	return L;
}


Spectrum ToFBipathIntegrator::LightPathL(PathData &data, const Sample &sample,
	ToFPathVertex* lightPath, const ToFPathVertex &vc, int *nSpecularVertices, 
	int cameraPathIdx, int lightPathIdx) const {
	const Point &pc = vc.bsdf->dgShading.p;
	const Normal &nc = vc.bsdf->dgShading.nn;
	const ToFPathVertex &vl = lightPath[lightPathIdx];
	const Point &pl = vl.bsdf->dgShading.p;
	const Normal &nl = vl.bsdf->dgShading.nn;
	Spectrum L(0.f);

	if (!vl.specularBounce) {
		// Compute contribution between camera and light vertices
		Vector w = Normalize(pl - pc);
		Spectrum fc = vc.bsdf->f(vc.wPrev, w) * (1 + vc.nSpecularComponents);
		Spectrum fl = vl.bsdf->f(-w, vl.wPrev) * (1 + vl.nSpecularComponents);
		Ray r(pc, pl - pc, 1e-3f, .999f, sample.time);

		if (!fc.IsBlack() && !fl.IsBlack() && !data.scene->IntersectP(r)) {
			// Compute weight for bidirectional path, _pathWt_
			int pathIdx = cameraPathIdx + lightPathIdx + 2;
			float pathWt = 1.f / (pathIdx - nSpecularVertices[pathIdx]);
			float G = AbsDot(nc, w) * AbsDot(nl, w) / DistanceSquared(pl, pc);
			L += (vc.alpha * fc * G * fl * vl.alpha) * pathWt;
		}
	}
	return L;
}


ToFBipathIntegrator *CreateToFBipathSurfaceIntegrator(
	const ParamSet &params) {
    int maxDepth = params.FindOneInt("maxdepth", 5);
	float lambdaMin = params.FindOneFloat("lambdamin", 400);
	float lambdaMax = params.FindOneFloat("lambdamax", 1000);
	return new ToFBipathIntegrator(maxDepth, lambdaMin, lambdaMax);
}
